from sage.combinat.backtrack import SearchForest
from sage.rings.integer import Integer
from sage.categories.finite_enumerated_sets import FiniteEnumeratedSets
from sage.structure.unique_representation import UniqueRepresentation

class RowCodes_n(UniqueRepresentation, SearchForest2):

    """
    The enumerated sets of row code of k-shape of length n

    This is a Catalan set

    EXAMPLES::

        sage: RowCodes_n(0).list()
        [[]]
        sage: RowCodes_n(1).list()
        [[0]]
        sage: RowCodes_n(2).list()
        [[0, 0], [1, 0]]
        sage: RowCodes_n(3).list()
        [[0, 0, 0], [1, 0, 0], [2, 0, 0], [0, 1, 0], [2, 1, 0]]
        sage: RowCodes_n(4).list()
        [[0, 0, 0, 0], [1, 0, 0, 0], [2, 0, 0, 0], [3, 0, 0, 0], [0, 1, 0, 0], [1, 1, 0, 0], [3, 1, 0, 0], [0, 2, 0, 0], [3, 2, 0, 0], [0, 0, 1, 0], [2, 0, 1, 0], [3, 0, 1, 0], [0, 2, 1, 0], [3, 2, 1, 0]]
    """
    def __init__(self, k):
        """
        TESTS::

            sage: for i in range(6): TestSuite(RowCodes_n(4)).run()
        """
        self.k = Integer(k)
        SearchForest2.__init__(self, reduce_init = [],category = FiniteEnumeratedSets())

    def roots(self):
        r"""
        The root of enumerations as per
        :class:`~sage.combinat.backtrack.SearchForest`

        EXAMPLES::

            sage: RowCodes_n(2).roots()
            [([], 0, [])]
        """
        return [([], 0, [])]

    def children(self, (code, length, prev_next)):
        """
        TESTS::

            sage: II = RowCodes_n(6)
            sage: x = II.roots()[0]
            sage: l = list(II.children(x)); l
            [([0], 1, [0])]
            sage: l = list(II.children(l[0])); l
            [([0, 0], 2, [0, 1]), ([1, 0], 2, [0])]
            sage: l = list(II.children(l[0])); l
            [([0, 0, 0], 3, [0, 1, 2]), ([1, 0, 0], 3, [0, 1]), ([2, 0, 0], 3, [0])]
            sage: l = list(II.children(l[2])); l
            [([0, 2, 0, 0], 4, [0, 3]), ([3, 2, 0, 0], 4, [0])]
            sage: l = list(II.children(l[0])); l
            [([0, 0, 2, 0, 0], 5, [0, 3, 4]), ([3, 0, 2, 0, 0], 5, [0, 1]), ([4, 0, 2, 0, 0], 5, [0])]
            sage: l = list(II.children(l[1])); l
            [([0, 3, 0, 2, 0, 0], 6, [0, 1, 5]), ([1, 3, 0, 2, 0, 0], 6, [0, 4]), ([5, 3, 0, 2, 0, 0], 6, [0])]
            sage: l = list(II.children(l[2])); l
            []
        """
        if length == self.k:
            return
        next = prev_next[:]
        next.append(length)
        length += 1
        for i in xrange(len(next)):
            nexti = next[i]
            yield ([nexti]+code, length, [j - nexti for j in next[i:]])

    def __contains__(self, x):
        """
        TESTS::

            sage: [] in RowCodes_n(0)
            True
            sage: [0,0] in RowCodes_n(2)
            True
            sage: [1,0] in RowCodes_n(2)
            True
            sage: [2,0] in RowCodes_n(2)
            False
            sage: [5, 3, 0, 2, 0, 0] in RowCodes_n(6)
            True
        """
        try:
            x = list(x)
        except TypeError:
            return False
        if len(x) != self.k: return False
        try:
            self.to_catalan_code(x)
        except ValueError:
            return False
        return True

    def to_catalan_code(self, x):
        """
        Compute the Catalan code associated to x

        A Catalan is a list of integer l such that `l_0 = 0` and
        `0 \leq l_{i-1} \leq l_i + 1` for `i > 0`

        EXAMPLES::

            sage: RowCodes_n(0).to_catalan_code([])
            []
            sage: RowCodes_n(2).to_catalan_code([0, 0])
            [0, 0]
            sage: RowCodes_n(2).to_catalan_code([1, 0])
            [0, 1]
            sage: RowCodes_n(6).to_catalan_code([5, 3, 0, 2, 0, 0])
            [0, 0, 2, 0, 1, 2]
        """
        res = []
        prefix = self.roots()[0]
        for i in range(1, len(x)+1):
            for j, new_prefix in enumerate(self.children(prefix)):
                if new_prefix[0] == x[-i:]:
                    prefix = new_prefix
                    res.append(j)
                    break
            else:
                raise ValueError, "%s is not a row-code"%x
        return res

    def post_process(self, (code, length, prev_next)):
        """
        Only consider the code of the given lenght

        TESTS::

            sage: I = RowCodes_n(0)
            sage: I.post_process(I.roots()[0])
            []
            sage: I = RowCodes_n(2)
            sage: I.post_process(I.roots()[0])
            sage: I.post_process(([0, 0], 2, []))
            [0, 0]
        """
        if length == self.k:
            return code

    def cardinality(self):
        """
        The row codes are counted by catalan numbers

        EXAMPLES::

            sage: [RowCodes_n(i).cardinality() for i in range(7)]
            [1, 1, 2, 5, 14, 42, 132]
        """
        from sage.combinat.combinat import catalan_number
        return catalan_number(self.k)

    def weight(self, x):
        """
        Return the weigth of the k-shape associated to x

        EXAMPLES::

            sage: I = RowCodes_n(2)
            sage: I.weight([0,0])
            0
            sage: I.weight([1,0])
            1
        """
        return sum((i+1)*j for i,j in enumerate(x))
	
    def reduce_function(self, x, y):
        return x + y
    
    def statistic(self, e):
        return [e]

